package cn.codingjc.aurora.framework.schedule.core;

import org.springframework.beans.factory.DisposableBean;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.config.CronTask;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @auth: coding_jc
 * @date: 2023/11/1 22:20
 * @desc:
 */
public class CronTaskRegistry implements DisposableBean {

    private final Map<Runnable, ScheduledTask> scheduledTaskMap = new ConcurrentHashMap<>(16);

    private TaskScheduler taskScheduler;

    public CronTaskRegistry(TaskScheduler taskScheduler) {
        this.taskScheduler = taskScheduler;
    }

    public TaskScheduler getTaskScheduler() {
        return taskScheduler;
    }

    public void addCronTask(Runnable runnable, String cronExpression){
        addCronTask(new CronTask(runnable, cronExpression));
    }

    public void addCronTask(CronTask cronTask){
        if (cronTask != null) {
            Runnable task = cronTask.getRunnable();
            if (scheduledTaskMap.containsKey(task)) {
                removeCronTask(task);
            }
            this.scheduledTaskMap.put(task, scheduleCronTask(cronTask));
        }
    }

    private void removeCronTask(Runnable task) {
        ScheduledTask scheduledTask = this.scheduledTaskMap.remove(task);
        if (scheduledTask != null) {
            scheduledTask.cancel();
        }
    }

    private ScheduledTask scheduleCronTask(CronTask cronTask) {
        ScheduledTask scheduledTask = new ScheduledTask();
        scheduledTask.future = this.taskScheduler.schedule(cronTask.getRunnable(), cronTask.getTrigger());
        return scheduledTask;
    }

    @Override
    public void destroy() throws Exception {
        for (ScheduledTask scheduledTask : this.scheduledTaskMap.values()) {
            scheduledTask.cancel();
        }
        this.scheduledTaskMap.clear();
    }
}
